# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""异步函数中的表达式测试

测试async/await语法中的复杂表达式，包括：
- for/for-each 循环
- 推导式表达式（列表推导式、集合推导式、字典推导式、生成器表达式）
- 字面量构造（列表、元组、字典、集合）
"""

from gtest import EXPECT_EQ, EXPECT_TRUE, test

from mc.future import ensure_future


# ============================================================================
# ensure_future 验证测试
# ============================================================================


@test("EnsureFuture", "BasicEnsureFuture")
def test_basic_ensure_future():
    """测试 ensure_future 基础功能"""

    async def async_func() -> int:
        return 42

    # 使用 ensure_future 包装异步函数并同步等待结果
    result = ensure_future(async_func()).get()
    EXPECT_EQ(result, 42)


@test("EnsureFuture", "EnsureFutureWithParameters")
def test_ensure_future_with_parameters():
    """测试 ensure_future 带参数的异步函数"""

    async def async_add(a: int, b: int) -> int:
        return a + b

    result = ensure_future(async_add(10, 20)).get()
    EXPECT_EQ(result, 30)


@test("EnsureFuture", "EnsureFutureNested")
def test_ensure_future_nested():
    """测试 ensure_future 嵌套调用"""

    async def inner() -> int:
        return 5

    async def outer() -> int:
        val = await inner()
        return val * 2

    result = ensure_future(outer()).get()
    EXPECT_EQ(result, 10)


# ============================================================================
# for-each 循环测试
# ============================================================================


@test("AsyncForEach", "ForEachList")
async def test_for_each_list():
    """测试 for-each 遍历列表"""

    async def get_value(x: int) -> int:
        return x + 10

    total = 0
    numbers: list[int] = [i for i in range(1, 101)]
    for num in numbers:
        val = await get_value(num)
        total = total + val

    EXPECT_EQ(total, 6050)


@test("AsyncForEach", "ListIterator")
async def test_list_iterator_with_continue():
    """测试手动迭代循环中的 continue"""

    async def get_value(x: int) -> int:
        return x + 10

    total = 0
    numbers = [1, 2, 3, 4, 5, 6]
    it = iter(numbers)
    while True:
        num = next(it, None)
        if not num:
            break
        # 跳过偶数
        if num % 2 == 0:
            continue
        val = await get_value(num)
        total = total + val

    # 只累加奇数: (1+10) + (3+10) + (5+10) = 11+13+15 = 39
    EXPECT_EQ(total, 39)


@test("AsyncForEach", "ForEachListWithBreak")
async def test_for_each_list_with_break():
    """测试 for-each 遍历列表中的 break"""

    async def get_value(x: int) -> int:
        return x * 2

    total = 0
    numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    for num in numbers:
        val = await get_value(num)
        if val > 10:
            break
        total = total + val

    # 1*2=2, 2*2=4, 3*2=6, 4*2=8, 5*2=10, 6*2=12 > 10 break
    # total = 2+4+6+8+10 = 30
    EXPECT_EQ(total, 30)


@test("AsyncForEach", "ForEachListWithContinue")
async def test_for_each_list_with_continue():
    """测试 for-each 遍历列表中的 continue"""

    async def get_value(x: int) -> int:
        return x

    total = 0
    numbers = [1, 2, 3, 4, 5]
    for num in numbers:
        val = await get_value(num)
        if val % 2 == 0:
            continue
        total = total + val

    # 跳过偶数：1, 3, 5 累加 = 9
    EXPECT_EQ(total, 9)


@test("AsyncForEach", "ForEachDict")
async def test_for_each_dict():
    """测试 for-each 遍历字典（遍历键）"""

    async def get_value(x: int) -> int:
        return x * 2

    total = 0
    d = {1: 10, 2: 20, 3: 30}
    for key in d:
        val = await get_value(key)
        total = total + val

    # 1*2 + 2*2 + 3*2 = 2+4+6 = 12
    EXPECT_EQ(total, 12)


@test("AsyncForEach", "ForEachDictItems")
async def test_for_each_dict_items():
    """测试 for-each 遍历字典的 items（元组解包）"""

    async def get_value(x: int) -> int:
        return x

    total = 0
    d = {1: 10, 2: 20, 3: 30}
    for key, value in d.items():
        val = await get_value(key)
        total = total + val + value

    # (1+10) + (2+20) + (3+30) = 11+22+33 = 66
    EXPECT_EQ(total, 66)


@test("AsyncForEach", "ForEachSet")
async def test_for_each_set():
    """测试 for-each 遍历集合"""

    async def get_value(x: int) -> int:
        return x + 5

    total = 0
    s = {1, 2, 3, 4, 5}
    for num in s:
        val = await get_value(num)
        total = total + val

    # 注意：集合遍历顺序不确定，但总和应该相同
    # (1+5) + (2+5) + (3+5) + (4+5) + (5+5) = 6+7+8+9+10 = 40
    EXPECT_EQ(total, 40)


@test("AsyncForEach", "NestedForEach")
async def test_nested_for_each():
    """测试嵌套的 for-each 循环"""

    async def get_value(x: int, y: int) -> int:
        return x + y

    total = 0
    matrix = [[1, 2], [3, 4], [5, 6]]
    for row in matrix:
        for num in row:
            val = await get_value(num, 10)
            total = total + val

    # (1+10) + (2+10) + (3+10) + (4+10) + (5+10) + (6+10) = 11+12+13+14+15+16 = 81
    EXPECT_EQ(total, 81)


@test("AsyncForEach", "ForEachWithAwaitInIterable")
async def test_for_each_with_await_in_iterable():
    """测试 for-each 的迭代对象中包含 await"""

    async def get_list() -> list[int]:
        return [1, 2, 3, 4, 5]

    async def get_value(x: int) -> int:
        return x * 2

    total = 0
    numbers = await get_list()
    for num in numbers:
        val = await get_value(num)
        total = total + val

    # (1*2) + (2*2) + (3*2) + (4*2) + (5*2) = 2+4+6+8+10 = 30
    EXPECT_EQ(total, 30)


# # ============================================================================
# # 列表推导式测试
# # ============================================================================


@test("AsyncListComp", "BasicListComp")
async def test_basic_list_comp():
    """测试基础列表推导式"""

    async def get_value(x: int) -> int:
        return x * 2

    numbers = [1, 2, 3, 4, 5]
    # TODO:: 类型推导引擎还不支持推导异步列表推导式的类型，需要手动指定 result 类型
    result: list[int] = [await get_value(x) for x in numbers]

    EXPECT_EQ(len(result), 5)
    EXPECT_EQ(result[0], 2)
    EXPECT_EQ(result[1], 4)
    EXPECT_EQ(result[2], 6)
    EXPECT_EQ(result[3], 8)
    EXPECT_EQ(result[4], 10)


@test("AsyncListComp", "ListCompWithCondition")
async def test_list_comp_with_condition():
    """测试带条件的列表推导式"""

    async def get_value(x: int) -> int:
        return x

    numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    result: list[int] = [await get_value(x) for x in numbers if x % 2 == 0]

    EXPECT_EQ(len(result), 5)
    EXPECT_EQ(result[0], 2)
    EXPECT_EQ(result[1], 4)
    EXPECT_EQ(result[2], 6)
    EXPECT_EQ(result[3], 8)
    EXPECT_EQ(result[4], 10)


@test("AsyncListComp", "ListCompWithAwaitInCondition")
async def test_list_comp_with_await_in_condition():
    """测试条件中包含 await 的列表推导式"""

    async def check_even(x: int) -> bool:
        return x % 2 == 0

    async def get_value(x: int) -> int:
        return x * 2

    numbers = [1, 2, 3, 4, 5]
    result: list[int] = [await get_value(x) for x in numbers if await check_even(x)]

    EXPECT_EQ(len(result), 2)
    EXPECT_EQ(result[0], 4)  # 2*2
    EXPECT_EQ(result[1], 8)  # 4*2


@test("AsyncListComp", "NestedListComp")
async def test_nested_list_comp():
    """测试嵌套列表推导式"""

    async def get_value(x: int, y: int) -> int:
        return x + y

    matrix = [[1, 2], [3, 4], [5, 6]]
    result: list[int] = [await get_value(x, 10) for row in matrix for x in row]

    EXPECT_EQ(len(result), 6)
    EXPECT_EQ(result[0], 11)  # 1+10
    EXPECT_EQ(result[1], 12)  # 2+10
    EXPECT_EQ(result[2], 13)  # 3+10
    EXPECT_EQ(result[3], 14)  # 4+10
    EXPECT_EQ(result[4], 15)  # 5+10
    EXPECT_EQ(result[5], 16)  # 6+10


@test("AsyncListComp", "ListCompWithAwaitInIterable")
async def test_list_comp_with_await_in_iterable():
    """测试迭代对象中包含 await 的列表推导式"""

    async def get_list() -> list[int]:
        return [1, 2, 3, 4, 5]

    async def get_value(x: int) -> int:
        return x * 2

    numbers = await get_list()
    result: list[int] = [await get_value(x) for x in numbers]

    EXPECT_EQ(len(result), 5)
    EXPECT_EQ(result[0], 2)
    EXPECT_EQ(result[1], 4)
    EXPECT_EQ(result[2], 6)
    EXPECT_EQ(result[3], 8)
    EXPECT_EQ(result[4], 10)


# # ============================================================================
# # 集合推导式测试
# # ============================================================================


@test("AsyncSetComp", "BasicSetComp")
async def test_basic_set_comp():
    """测试基础集合推导式"""

    async def get_value(x: int) -> int:
        return x * 2

    numbers = [1, 2, 3, 4, 5]
    result: set[int] = {await get_value(x) for x in numbers}

    EXPECT_EQ(len(result), 5)
    EXPECT_TRUE(2 in result)
    EXPECT_TRUE(4 in result)
    EXPECT_TRUE(6 in result)
    EXPECT_TRUE(8 in result)
    EXPECT_TRUE(10 in result)


@test("AsyncSetComp", "SetCompWithCondition")
async def test_set_comp_with_condition():
    """测试带条件的集合推导式"""

    async def get_value(x: int) -> int:
        return x

    numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    result: set[int] = {await get_value(x) for x in numbers if x % 2 == 0}

    EXPECT_EQ(len(result), 5)
    EXPECT_TRUE(2 in result)
    EXPECT_TRUE(4 in result)
    EXPECT_TRUE(6 in result)
    EXPECT_TRUE(8 in result)
    EXPECT_TRUE(10 in result)


@test("AsyncSetComp", "SetCompWithAwaitInCondition")
async def test_set_comp_with_await_in_condition():
    """测试条件中包含 await 的集合推导式"""

    async def check_even(x: int) -> bool:
        return x % 2 == 0

    async def get_value(x: int) -> int:
        return x * 2

    numbers = [1, 2, 3, 4, 5]
    result: set[int] = {await get_value(x) for x in numbers if await check_even(x)}

    EXPECT_EQ(len(result), 2)
    EXPECT_TRUE(4 in result)  # 2*2
    EXPECT_TRUE(8 in result)  # 4*2


# # ============================================================================
# # 字典推导式测试
# # ============================================================================


@test("AsyncDictComp", "BasicDictComp")
async def test_basic_dict_comp():
    """测试基础字典推导式"""

    async def get_value(x: int) -> int:
        return x * 2

    numbers = [1, 2, 3, 4, 5]
    result: dict[int, int] = {x: await get_value(x) for x in numbers}

    EXPECT_EQ(len(result), 5)
    EXPECT_EQ(result[1], 2)
    EXPECT_EQ(result[2], 4)
    EXPECT_EQ(result[3], 6)
    EXPECT_EQ(result[4], 8)
    EXPECT_EQ(result[5], 10)


@test("AsyncDictComp", "DictCompWithCondition")
async def test_dict_comp_with_condition():
    """测试带条件的字典推导式"""

    async def get_value(x: int) -> int:
        return x * 2

    numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    result: dict[int, int] = {x: await get_value(x) for x in numbers if x % 2 == 0}

    EXPECT_EQ(len(result), 5)
    EXPECT_EQ(result[2], 4)
    EXPECT_EQ(result[4], 8)
    EXPECT_EQ(result[6], 12)
    EXPECT_EQ(result[8], 16)
    EXPECT_EQ(result[10], 20)


@test("AsyncDictComp", "DictCompWithAwaitInKeyValue")
async def test_dict_comp_with_await_in_key_value():
    """测试键和值中都包含 await 的字典推导式"""

    async def get_key(x: int) -> int:
        return x * 2

    async def get_value(x: int) -> int:
        return x * 3

    numbers = [1, 2, 3]
    result: dict[int, int] = {await get_key(x): await get_value(x) for x in numbers}

    EXPECT_EQ(len(result), 3)
    EXPECT_EQ(result[2], 3)  # 1*2: 1*3
    EXPECT_EQ(result[4], 6)  # 2*2: 2*3
    EXPECT_EQ(result[6], 9)  # 3*2: 3*3


# # ============================================================================
# # 字面量构造测试
# # ============================================================================


@test("AsyncLiterals", "ListLiteral")
async def test_list_literal():
    """测试列表字面量构造"""

    async def get_value(x: int) -> int:
        return x * 2

    result: list[int] = [await get_value(1), await get_value(2), await get_value(3)]

    EXPECT_EQ(len(result), 3)
    EXPECT_EQ(result[0], 2)
    EXPECT_EQ(result[1], 4)
    EXPECT_EQ(result[2], 6)


@test("AsyncLiterals", "TupleLiteral")
async def test_tuple_literal():
    """测试元组字面量构造"""

    async def get_value(x: int) -> int:
        return x * 2

    result: tuple[int, int, int] = (
        await get_value(1),
        await get_value(2),
        await get_value(3),
    )

    EXPECT_EQ(len(result), 3)
    EXPECT_EQ(result[0], 2)
    EXPECT_EQ(result[1], 4)
    EXPECT_EQ(result[2], 6)


@test("AsyncLiterals", "DictLiteral")
async def test_dict_literal():
    """测试字典字面量构造"""

    async def get_key(x: int) -> int:
        return x

    async def get_value(x: int) -> int:
        return x * 2

    result: dict[int, int] = {
        await get_key(1): await get_value(1),
        await get_key(2): await get_value(2),
        await get_key(3): await get_value(3),
    }

    EXPECT_EQ(len(result), 3)
    EXPECT_EQ(result[1], 2)
    EXPECT_EQ(result[2], 4)
    EXPECT_EQ(result[3], 6)


@test("AsyncLiterals", "SetLiteral")
async def test_set_literal():
    """测试集合字面量构造"""

    async def get_value(x: int) -> int:
        return x * 2

    result: set[int] = {await get_value(1), await get_value(2), await get_value(3)}

    EXPECT_EQ(len(result), 3)
    EXPECT_TRUE(2 in result)
    EXPECT_TRUE(4 in result)
    EXPECT_TRUE(6 in result)


@test("AsyncLiterals", "NestedListLiteral")
async def test_nested_list_literal():
    """测试嵌套列表字面量构造"""

    async def get_value(x: int) -> int:
        return x

    result: list[list[int]] = [
        list[int]([await get_value(1), await get_value(2)]),
        list[int]([await get_value(3), await get_value(4)]),
    ]

    EXPECT_EQ(len(result), 2)
    EXPECT_EQ(len(result[0]), 2)
    EXPECT_EQ(len(result[1]), 2)
    EXPECT_EQ(result[0][0], 1)
    EXPECT_EQ(result[0][1], 2)
    EXPECT_EQ(result[1][0], 3)
    EXPECT_EQ(result[1][1], 4)


@test("AsyncLiterals", "MixedLiteral")
async def test_mixed_literal():
    """测试混合字面量构造（列表包含字典等）"""

    async def get_key(x: int) -> int:
        return x

    async def get_value(x: int) -> int:
        return x * 2

    result: list[dict[int, int]] = [
        dict[int, int]({await get_key(1): await get_value(1)}),
        dict[int, int]({await get_key(2): await get_value(2)}),
    ]

    EXPECT_EQ(len(result), 2)
    EXPECT_EQ(len(result[0]), 1)
    EXPECT_EQ(len(result[1]), 1)
    EXPECT_EQ(result[0][1], 2)
    EXPECT_EQ(result[1][2], 4)


# # ============================================================================
# # 混合场景测试
# # ============================================================================


@test("AsyncMixed", "ForEachWithListComp")
async def test_for_each_with_list_comp():
    """测试 for-each 循环中使用列表推导式"""

    async def get_value(x: int) -> int:
        return x + 10

    numbers = [1, 2, 3]
    doubled = [await get_value(x) for x in numbers]
    total = 0
    for val in doubled:
        total = total + val

    # (1+10) + (2+10) + (3+10) = 11+12+13 = 36
    EXPECT_EQ(total, 36)


@test("AsyncMixed", "ListCompWithLiteral")
async def test_list_comp_with_literal():
    """测试列表推导式中使用字面量"""

    async def get_value(x: int) -> int:
        return x * 2

    numbers = [1, 2, 3]
    a: list[int] = [await get_value(x) for x in numbers]
    b = list[int](
        [
            await get_value(10),
            await get_value(20),
        ]
    )
    result = a + b

    EXPECT_EQ(len(result), 5)
    EXPECT_EQ(result[0], 2)
    EXPECT_EQ(result[1], 4)
    EXPECT_EQ(result[2], 6)
    EXPECT_EQ(result[3], 20)
    EXPECT_EQ(result[4], 40)


@test("AsyncMixed", "NestedComprehensionWithAwait")
async def test_nested_comprehension_with_await():
    """测试嵌套推导式中包含 await"""

    async def get_value(x: int, y: int) -> int:
        return x + y

    matrix = [[1, 2], [3, 4]]
    result: list[int] = [
        await get_value(x, y) for row in matrix for x in row for y in [10, 20]
    ]

    # 1+10, 1+20, 2+10, 2+20, 3+10, 3+20, 4+10, 4+20
    EXPECT_EQ(len(result), 8)
    EXPECT_EQ(result[0], 11)
    EXPECT_EQ(result[1], 21)
    EXPECT_EQ(result[2], 12)
    EXPECT_EQ(result[3], 22)
    EXPECT_EQ(result[4], 13)
    EXPECT_EQ(result[5], 23)
    EXPECT_EQ(result[6], 14)
    EXPECT_EQ(result[7], 24)
